"""
https://www.tensorflow.org/guide/keras/train_and_evaluate
"""

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, losses, metrics, optimizers, Model
from python_ai.common.xcommon import *
import os
import sys
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

# data
x, y = load_breast_cancer(return_X_y=True)
x = StandardScaler().fit_transform(x)

x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=777)
m, n = x_train.shape

# Reserve 10,000 samples for validation
n_val = int(np.ceil(m * 0.1))
x_val = x_train[-n_val:]
y_val = y_train[-n_val:]
x_train = x_train[:-n_val]
y_train = y_train[:-n_val]

# model
inputs = keras.Input(shape=(n,), name="digits")
x = layers.Dense(64, activation="relu", name="dense_1")(inputs)
x = layers.Dense(64, activation="relu", name="dense_2")(x)
outputs = layers.Dense(1, activation="sigmoid", name="predictions")(x)

model = Model(inputs=inputs, outputs=outputs)


class MyMetric(metrics.Metric):

    def __init__(self, name='acc', **kwargs):
        super(MyMetric, self).__init__(name=name, **kwargs)
        self.n_total = tf.Variable(0, dtype=tf.int32, name='n_total')
        self.n_equal = tf.Variable(0, dtype=tf.int32, name='n_total')

    def reset_state(self):
        self.n_total.assign(0)
        self.n_equal.assign(0)

    def result(self):
        return tf.cast(self.n_equal, tf.float32) / tf.cast(self.n_total, tf.float32)

    def update_state(self, y_true, y_pred, sample_weight=None):
        self.n_total.assign_add(len(y_true))
        self.n_equal.assign_add(tf.reduce_sum(tf.cast(tf.equal(
            tf.cast(y_true, tf.float32) > 0.5,
            y_pred > 0.5
        ), tf.int32)))


# specify the training config
model.compile(
    optimizer=optimizers.RMSprop(),
    loss=losses.binary_crossentropy,
    metrics=[MyMetric('my_acc'), metrics.binary_accuracy]
)

# fit and train
sep('Fit and train')
history = model.fit(
    x_train,
    y_train,
    batch_size=64,
    epochs=2,
    validation_data=(x_val, y_val)
)
print(history.history)

sep('Evaluate on test data')
results = model.evaluate(x_test, y_test, batch_size=128)
print(f'test loss, test acc: {results}')

